/*
 * Copyright 2020 Makani Technologies LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef LIB_DATATOOLS_LOAD_PCAP_HDF5_DATASET_H_
#define LIB_DATATOOLS_LOAD_PCAP_HDF5_DATASET_H_

#include <glog/logging.h>
#include <hdf5.h>
#include <hdf5_hl.h>
#include <stddef.h>
#include <stdint.h>

#include <functional>
#include <string>
#include <vector>

// Provide access to the raw bytes of a dataset contained in an HDF5
// file generated by pcap_to_hdf5.
class PcapToHdf5DatasetRawData {
 public:
  // Constructor for PcapToHdf5DatasetRawData.
  //
  // Args:
  //   filename: Path to and HDF5 file.
  //   aio_node_str: See AioNodeToString.
  //   message_type_str: See MessageTypeToString.
  //   first_message_to_read: First message to read from the file.
  //   last_message_to_read: Last message to read or -1 indicating
  //       all messages should be read.
  PcapToHdf5DatasetRawData(const std::string &filename,
                           const std::string &aio_node_str,
                           const std::string &message_type_str,
                           int32_t first_message_to_read,
                           int32_t last_message_to_read);

  // Returns the number of messages read.
  int32_t GetNumMessages() const { return num_messages_; }

  // Fetches a pointer to the first byte of a given structure containing
  // a message and header information (see pcap_to_hdf5).
  const uint8_t *GetMessageData(int32_t i) const {
    CHECK_LE(0, i);
    CHECK_LT(i, num_messages_);
    return message_and_header_raw_data_.data()
        + i * message_and_header_size_ + message_offset_;
  }

  // Returns the size of the message portion of each dataset element.
  size_t GetMessageSize() const { return message_size_; }

 private:
  // Helper function which populates this objects member variables once
  // the HDF5 dataset has been identified.
  void LoadFromDataset(hid_t message_and_header_did,
                       int32_t first_message_to_read,
                       int32_t last_message_to_read);

  // Number of messages read.
  int32_t num_messages_;

  // Size (in bytes) of a combined message and header structure.
  size_t message_and_header_size_;

  // Size of the message portion of each dataset element.
  size_t message_size_;

  // Offset of the message portion of the dataset element.
  size_t message_offset_;

  // All the raw data.
  std::vector<uint8_t> message_and_header_raw_data_;
};

// Class for accessing and unpacking data from an HDF5 dataset generated
// by pcap_to_hdf5.
template <typename T>
class PcapToHdf5Dataset {
 public:
  // Constructor.
  //
  // Args:
  //   filename: Path to and HDF5 file.
  //   aio_node_str: See AioNodeToString.
  //   message_type_str: See MessageTypeToString.
  //   first_message_to_read: First message to read from the file.
  //   last_message_to_read: Last message to read or -1 indicating
  //       all messages should be read.
  //   packed_size: Size of T when packed (e.g. PACK_CONTROLTELEMETRY_SIZE).
  //   unpack: Unpacking function (e.g. UnpackControlTelemetry).
  PcapToHdf5Dataset(
      const std::string &filename, const std::string &aio_node_str,
      const std::string &message_type_str, int32_t first_message_to_read,
      int32_t last_message_to_read, size_t packed_size,
      const std::function<size_t(const uint8_t *, size_t, T *)> &unpack);
  ~PcapToHdf5Dataset() {}

  // Returns the number of messages in this dataset.
  size_t GetSize() const { return raw_data_.GetNumMessages(); }

  // Unpacks a specific datum in to a given message.
  void UnpackMessage(int32_t i, T *message) const;

 private:
  const std::function<size_t(const uint8_t *, size_t, T *)> unpack_;

  // Raw bytes associated with the dataset.
  PcapToHdf5DatasetRawData raw_data_;
};

template <typename T>
PcapToHdf5Dataset<T>::PcapToHdf5Dataset(
    const std::string &filename,
    const std::string &aio_node_str,
    const std::string &message_type_str,
    int32_t first_message_to_read, int32_t last_message_to_read,
    size_t packed_size,
    const std::function<size_t(const uint8_t *, size_t, T *)> &unpack)
    : unpack_(unpack),
      raw_data_(filename, aio_node_str, message_type_str, first_message_to_read,
                last_message_to_read) {
  CHECK_EQ(packed_size, raw_data_.GetMessageSize());
}

template <typename T>
void PcapToHdf5Dataset<T>::UnpackMessage(int32_t i, T *message) const {
  unpack_(raw_data_.GetMessageData(i), 1, message);
}

// Open a dataset and get the parameter requested.
//
// This is returned through a user-provided pointer whose type's size is checked
// against the data size in the log file.
//
// Args:
//   filename: Path to a HDF5 file.
//   param_type_str: Which parameter to load. For example:
//       "sim_params", "system_params" or "control_params".
//   result: Pointer to struct where loaded message will be written.
template <typename T>
void LoadHdf5Parameters(
    const std::string &filename, const std::string &param_type_str, T *result) {
  // Open the dataset containing the parameter given by param_type_str.
  // This involves opening the HDF5 file and the parameter group, all released
  // at the end of this function.
  CHECK_NOTNULL(result);

  // Open the file.
  hid_t file_id = H5Fopen(filename.c_str(), H5F_ACC_RDONLY, H5P_DEFAULT);
  CHECK_LE(0, file_id) << "Could not open file.";

  // Get the parameters group.
  hid_t parameters_group_id = H5Gopen2(file_id, "parameters", H5P_DEFAULT);
  CHECK_LE(0, parameters_group_id) << "Could not open /parameters.";

  // Get the requested parameter's dataset ID.
  hid_t param_dataset_id =
      H5Dopen2(parameters_group_id, param_type_str.c_str(), H5P_DEFAULT);
  CHECK_LE(0, param_dataset_id)
      << "Could not open " << param_type_str << ".";

  // Get the requested parameter's type ID.
  hid_t param_type_id = H5Dget_type(param_dataset_id);
  CHECK_LE(0, param_type_id) << "Could not get parameter type.";

  // Get the requested parameter's size.
  size_t param_size = H5Tget_size(param_type_id);
  CHECK_EQ(sizeof(*result), param_size)
      << "Return type size doesn't match log size.";

  // Read the requested parameter into the result pointer.
  CHECK_LE(0, H5Dread(param_dataset_id, param_type_id, H5S_ALL, H5S_ALL,
                      H5P_DEFAULT, result));

  // Tidy up the memory.
  H5Tclose(param_type_id);
  H5Dclose(param_dataset_id);
  H5Gclose(parameters_group_id);
  H5Fclose(file_id);
}

#endif  // LIB_DATATOOLS_LOAD_PCAP_HDF5_DATASET_H_
